#include "data_utils.h"
#include "kernel_tiling/kernel_tiling.h"
#include "tiling/platform/platform_ascendc.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_transpose_custom.h"
#include "aclrtlaunch_reducesum_custom.h"

#else
#include "acl/acl.h"
#include "aclrtlaunch_transpose_custom.h"
#include "tikicpulib.h"
extern "C" void transpose_custom(uint8_t *, uint8_t *);
#endif

void dumpToTextFile(void *scratchHost,
                    size_t dstFileSize,
                    int rows,
                    int cols,
                    const char *filename)
{
    uint16_t *float16Data = reinterpret_cast<uint16_t *>(scratchHost);

    // 1) 构造带 .txt 后缀的完整文件名
    char filepath[512];
    int ret = snprintf(filepath, sizeof(filepath), "./data/%s.txt", filename);
    if (ret < 0 || ret >= (int)sizeof(filepath))
    {
        fprintf(stderr, "文件名过长或格式化出错\n");
        return;
    }

    // 2) 打开文件
    FILE *fp = fopen(filepath, "w");
    if (!fp)
    {
        perror("fopen");
        return;
    }

    // 3) 按矩阵格式写入
    for (int r = 0; r < rows; ++r)
    {
        for (int c = 0; c < cols; ++c)
        {
            int idx = r * cols + c;
            float f = aclFloat16ToFloat(float16Data[idx]);
            fprintf(fp, "%8.3f", f);
            if (c + 1 < cols)
                fputc(' ', fp);
        }
        fputc('\n', fp);
    }

    fclose(fp);
    printf("已将数据按 %d×%d 矩阵格式写入 %s\n", rows, cols, filepath);
}

void dumpToTextFilefp32(void *scratchHost,
                        size_t dstFileSize,
                        int rows,
                        int cols,
                        const char *filename)
{
    float *float32Data = reinterpret_cast<float *>(scratchHost);

    // 1) 构造带 .txt 后缀的完整文件名
    char filepath[512];
    int ret = snprintf(filepath, sizeof(filepath), "./data/%s.txt", filename);
    if (ret < 0 || ret >= (int)sizeof(filepath))
    {
        fprintf(stderr, "文件名过长或格式化出错\n");
        return;
    }

    // 2) 打开文件
    FILE *fp = fopen(filepath, "w");
    if (!fp)
    {
        perror("fopen");
        return;
    }

    // 3) 按矩阵格式写入
    for (int r = 0; r < rows; ++r)
    {
        for (int c = 0; c < cols; ++c)
        {
            int idx = r * cols + c;
            float f = float32Data[idx];
            fprintf(fp, "%8.3f", f);
            if (c + 1 < cols)
                fputc(' ', fp);
        }
        fputc('\n', fp);
    }

    fclose(fp);
    printf("已将数据按 %d×%d 矩阵格式写入 %s\n", rows, cols, filepath);
}

int main(int32_t argc, char *argv[])
{
    // const char *socVersion = SOC_VERSION;
    // auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(socVersion);
    size_t srcFileSize = 256 * 256 * 4 * sizeof(int16_t);
    size_t dstFileSize = 256 * 256 * 4 * sizeof(int16_t);

    uint32_t blockDim = 32;

    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 2;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));
    uint8_t *srcHost;
    uint8_t *srcDevice;
    CHECK_ACL(aclrtMallocHost((void **)(&srcHost), srcFileSize));
    CHECK_ACL(aclrtMalloc((void **)&srcDevice, srcFileSize, ACL_MEM_MALLOC_HUGE_FIRST));
    ReadFile("./input/x1_gm.bin", srcFileSize, srcHost, srcFileSize);

    printf("the srcHost is :\n");
    uint16_t *float16Data = reinterpret_cast<uint16_t *>(srcHost);
    int elementsPerRow = 16; // 假设 K 是你的列数
    int flag = 0;

    for (int i = 0; i < 16 * 16; i++)
    {
        if (flag == elementsPerRow)
        {
            printf("\n");
            flag = 0;
        }
        printf(" %f", aclFloat16ToFloat(float16Data[i]));
        flag++;
    }
    printf("\n");

    CHECK_ACL(aclrtMemcpy(srcDevice, srcFileSize, srcHost, srcFileSize, ACL_MEMCPY_HOST_TO_DEVICE));

    // CHECK_ACL(aclrtMemcpy(srcHost, srcFileSize, srcDevice, srcFileSize, ACL_MEMCPY_DEVICE_TO_HOST));

    printf("查看device的接收情况\n");
    float16Data = reinterpret_cast<uint16_t *>(srcHost);
    flag = 0;

    for (int i = 0; i < 16 * 16; i++)
    {
        if (flag == elementsPerRow)
        {
            printf("\n");
            flag = 0;
        }
        printf(" %f", aclFloat16ToFloat(float16Data[i]));
        flag++;
    }
    printf("\n");

    uint8_t *dstHost;
    uint8_t *scratchHost;
    uint8_t *scratchHost_reducesum;
    uint8_t *dstDevice;
    uint8_t *dstDevice_reducesum;
    CHECK_ACL(aclrtMallocHost((void **)(&dstHost), dstFileSize));
    CHECK_ACL(aclrtMallocHost((void **)(&scratchHost), dstFileSize));
    CHECK_ACL(aclrtMallocHost((void **)(&scratchHost_reducesum), dstFileSize));
    CHECK_ACL(aclrtMalloc((void **)&dstDevice, dstFileSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&dstDevice_reducesum, dstFileSize, ACL_MEM_MALLOC_HUGE_FIRST));

    uint8_t *scratch;
    CHECK_ACL(aclrtMalloc((void **)&scratch, dstFileSize, ACL_MEM_MALLOC_HUGE_FIRST));

    uint8_t *scratch_reducesum;
    CHECK_ACL(aclrtMalloc((void **)&scratch_reducesum, dstFileSize, ACL_MEM_MALLOC_HUGE_FIRST));

    dumpToTextFile(srcHost, dstFileSize, 512, 512, "src");

    auto ret = ACLRT_LAUNCH_KERNEL(transpose_custom)(blockDim, stream, srcDevice, dstDevice, scratch);
    CHECK_ACL(aclrtSynchronizeStream(stream));
    auto ret2 = ACLRT_LAUNCH_KERNEL(reducesum_custom)(blockDim, stream, srcDevice, dstDevice_reducesum, scratch_reducesum);

    CHECK_ACL(aclrtSynchronizeStream(stream));
    if (ret != ACL_SUCCESS && ret2 != ACL_SUCCESS)
    {
        const char *errMsg = aclGetRecentErrMsg();
        printf("Kernel execution error: %s\n", errMsg);
    }
    else
    {
        printf("Kernel execution %d\n", ret);
    }

    std::cout << "Copying results back to host..." << std::endl;
    CHECK_ACL(aclrtMemcpy(scratchHost, dstFileSize, scratch, dstFileSize, ACL_MEMCPY_DEVICE_TO_HOST));

    std::cout << "Copying scratch results back to host..." << std::endl;
    CHECK_ACL(aclrtMemcpy(dstHost, dstFileSize, dstDevice, dstFileSize, ACL_MEMCPY_DEVICE_TO_HOST));

    std::cout << "Copying reducesum  results back to host..." << std::endl;
    CHECK_ACL(aclrtMemcpy(scratchHost_reducesum, dstFileSize, scratch_reducesum, dstFileSize, ACL_MEMCPY_DEVICE_TO_HOST));

    dumpToTextFilefp32(scratchHost_reducesum, dstFileSize, 512, 512, "reducesum");

    dumpToTextFile(scratchHost, dstFileSize, 512, 512, "scratch");

    printf("the scratch is :\n");
    float16Data = reinterpret_cast<uint16_t *>(scratchHost);
    elementsPerRow = 16; // 假设 K 是你的列数
    flag = 0;
    for (int i = 0; i < 16 * 16; i++)
    { // 这里的32是指32个float16元素，不是字节
        if (flag == elementsPerRow)
        {
            printf("\n");
            flag = 0;
        }
        printf(" %f", aclFloat16ToFloat(float16Data[i]));
        flag++;
    }
    printf("\n");

    printf("the dstHost is :\n");
    float16Data = reinterpret_cast<uint16_t *>(dstHost);
    elementsPerRow = 16; // 假设 K 是你的列数
    flag = 0;

    for (int i = 0; i < 16 * 16; i++)
    { // 这里的32是指32个float16元素，不是字节
        if (flag == elementsPerRow)
        {
            printf("\n");
            flag = 0;
        }
        printf(" %f", aclFloat16ToFloat(float16Data[i]));
        flag++;
    }
    printf("\n");

    std::cout << "Cleaning up resources..." << std::endl;
    CHECK_ACL(aclrtFree(srcDevice));
    CHECK_ACL(aclrtFree(dstDevice));
    CHECK_ACL(aclrtFree(scratch));
    CHECK_ACL(aclrtFreeHost(srcHost));
    CHECK_ACL(aclrtFreeHost(dstHost));
    CHECK_ACL(aclrtFreeHost(scratchHost));
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());

    std::cout << "Program completed successfully!" << std::endl;
    return 0;
}